import matplotlib.pyplot as plt
import torch


def ave_loss_in_epoch(loss_array_in_epoch):
    with torch.no_grad():
        # 使用 torch.stack() 拼接张量，并保留在 GPU 上
        loss_tensor = torch.stack(loss_array_in_epoch)
        # 直接返回张量的均值
        return loss_tensor.mean().item()

def ave_score_in_epoch(score_dic_array_in_epoch):
    # 初始化
    score_dic_epoch_ave = {}
    for k, v in score_dic_array_in_epoch[0].items():
        score_dic_epoch_ave[k] = 0
    # 求平均
    iteration_num = len(score_dic_array_in_epoch)
    for i in range(iteration_num):
        for k, v in score_dic_array_in_epoch[i].items():
            score_dic_epoch_ave[k] += v/iteration_num

    # 求各项之和
    score_val_epoch_ave = 0
    for k, v in score_dic_epoch_ave.items():
        score_val_epoch_ave += v

    score_val_epoch_ave/=3
    return score_val_epoch_ave, score_dic_epoch_ave

def set_loss_plt_size(width, height):
    # 创建一个包含两个子图的图形
    fig, axs = plt.subplots(1, 2, figsize=(width, height))
    return fig , axs



def plot_loss(losses):
    # 绘制损失曲线
    plt.clf()
    plt.plot(losses, linewidth=2, color='firebrick')
    plt.tick_params(axis='both', size=5, width=2,
                    direction='in', labelsize=15)
    plt.xlabel('Epoch', size=15)
    plt.ylabel('Loss', size=15)
    plt.title('Training Loss Curve', size=20)
    plt.grid(color='midnightblue', linestyle='-.', linewidth=0.5)
    # 调整坐标轴的边框样式
    for spine in plt.gca().spines.values():
        spine.set_linewidth(2)  # 设置边框宽度为2
    plt.pause(0.0001)

def plot_loss_and_fid(fig, axs, loss_curve, fid_curve,epoch_curve):
    # 绘制第一个子图的损失曲线
    axs[0].cla()
    axs[0].plot(loss_curve, linewidth=2, color='firebrick')
    axs[0].tick_params(axis='both', size=5, width=2,
                       direction='in', labelsize=15)
    axs[0].set_xlabel('Epoch', size=15)
    axs[0].set_ylabel('Loss', size=15)
    axs[0].set_title('Training Loss Curve', size=20)
    axs[0].grid(color='midnightblue', linestyle='-.', linewidth=0.5)
    # 调整坐标轴的边框样式
    for spine in axs[0].spines.values():
        spine.set_linewidth(2)  # 设置边框宽度为2

    # 绘制第二个子图的损失曲线
    axs[1].cla()
    axs[1].plot(epoch_curve, fid_curve, linewidth=2, color='blue')
    axs[1].tick_params(axis='both', size=5, width=2,
                       direction='in', labelsize=15)
    axs[1].set_xlabel('Epoch', size=15)
    axs[1].set_ylabel('FID', size=15)
    axs[1].set_title('Training FID Curve', size=20)
    axs[1].grid(color='midnightblue', linestyle='-.', linewidth=0.5)
    # 调整坐标轴的边框样式
    for spine in axs[1].spines.values():
        spine.set_linewidth(2)  # 设置边框宽度为2
    plt.tight_layout()
    plt.pause(0.0001)



def plot_loss_and_score(fig, axs, loss_curve, sdf_scores = None, eikonal_scores = None):
    # 绘制第一个子图的损失曲线
    axs[0].cla()
    axs[0].plot(loss_curve, linewidth=2, color='firebrick')
    axs[0].tick_params(axis='both', size=5, width=2,
                       direction='in', labelsize=15)
    axs[0].set_xlabel('Epoch', size=15)
    axs[0].set_ylabel('Loss', size=15)
    axs[0].set_title('Training Loss Curve', size=20)
    axs[0].grid(color='midnightblue', linestyle='-.', linewidth=0.5)
    # 调整坐标轴的边框样式
    for spine in axs[0].spines.values():
        spine.set_linewidth(2)  # 设置边框宽度为2
    if hasattr(sdf_scores, "any"):
        # 绘制第二个子图的损失曲线
        axs[1].cla()
        # 绘制蓝色点
        axs[1].plot(sdf_scores, eikonal_scores, linewidth=2, marker='o', color='blue')
        # 绘制红色点
        axs[1].scatter(sdf_scores[-1], eikonal_scores[-1], marker='o', color='white', linewidth=2, edgecolors='red',s= 100, zorder=10)
        axs[1].tick_params(axis='both', size=5, width=2,
                        direction='in', labelsize=15)
        axs[1].set_xlabel('SDF Score', size=15)
        axs[1].set_ylabel('Eikonal Score', size=15)
        axs[1].set_title('Score Map', size=20)
        axs[1].grid(color='midnightblue', linestyle='-.', linewidth=0.5)
    # 调整坐标轴的边框样式
    for spine in axs[1].spines.values():
        spine.set_linewidth(2)  # 设置边框宽度为2
    plt.tight_layout()
    plt.pause(0.0001)


def save_plot(file_name, model_save_path):
    plt.savefig(f"{model_save_path}/{file_name}.png")

def save_curve(file_name, model_save_path, epoch, losses, create_new):
    # 如果需要新建，则清空文件，写入表头
    if create_new:
        if not hasattr(save_curve, file_name):
            with open(f"{model_save_path}/{file_name}.csv", 'w') as f:
                f.write("epoch, loss\n") #写入表头
            setattr(save_curve, file_name, True)
    # 追加保存曲线内容
    with open(f"{model_save_path}/{file_name}.csv", 'a') as f:
        f.write(f"{epoch}, {losses[-1]}\n") #追加数据

def save_map(file_name, model_save_path, epoch, dim_x, dim_y, create_new):
    # 如果需要新建，则清空文件，写入表头
    if create_new:
        if not hasattr(save_map, file_name):
            with open(f"{model_save_path}/{file_name}.csv", 'w') as f:
                f.write("epoch, sdf, eikonal\n") #写入表头
            setattr(save_map, file_name, True)
    # 追加保存曲线内容
    with open(f"{model_save_path}/{file_name}.csv", 'a') as f:
        f.write(f"{epoch}, {dim_x[-1]}, {dim_y[-1]}\n") #追加数据

def save_dictionary(file_name, model_save_path, epoch, dictionarys, create_new):
    # 如果需要新建，则清空文件，写入表头
    if create_new:
        if not hasattr(save_dictionary, file_name):
            with open(f"{model_save_path}/{file_name}.csv", 'w') as f:
                f.write("epoch") #写入表头
                for k, v in dictionarys[0].items():
                    f.write(f", {k}")
                f.write("\n")
            setattr(save_dictionary, file_name, True)
    # 追加保存字典内容
    with open(f"{model_save_path}/{file_name}.csv", 'a') as f:
        f.write(f"{epoch}")
        for k, v in dictionarys[-1].items():
            f.write(f", {v}")#追加数据
        f.write("\n")
